# -*- coding: utf-8 -*-
import cv2 as cv
from mediapipe.python.solutions import drawing_utils, objectron

cap = cv.VideoCapture("videos/shoe1.mp4")

mpObjectron = objectron.Objectron(static_image_mode=False,
                            max_num_objects=5,
                            min_detection_confidence=0.5,
                            min_tracking_confidence=0.99,
                            model_name='Shoe')

while True:
    flag, frame = cap.read()

    frameRGB = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
    results = mpObjectron.process(frameRGB)

    # Draw the box landmarks on the image.
    if results.detected_objects:
        for detected_object in results.detected_objects:
            drawing_utils.draw_landmarks(frame, detected_object.landmarks_2d, objectron.BOX_CONNECTIONS)
            drawing_utils.draw_axis(frame, detected_object.rotation, detected_object.translation)
    
    cv.imshow('MediaPipe Objectron', frame)
    if cv.waitKey(5) & 0xFF == ord('q'):
        break

cap.release()
cv.destroyAllWindows()